#include "kernel_operator.h"
#include "op_common.h"
namespace AscendC {

template <typename _DT_Y> class Kernel {
public:
  int64_t b = 0;
  int64_t r = 0;
  int64_t c = 0;
  DefGlobalTensor(Y);
public:
  __aicore__ inline Kernel() {}

  template <typename T> __aicore__ inline void Init(GM_ADDR y, GM_ADDR y_ref, T tiling_data) {
    this->b = tiling_data.b;
    this->r = tiling_data.r;
    this->c = tiling_data.c;
    // 切分global
    GSetBuffer(Y, y, 0, b * r * c);
  }
  __aicore__ inline void Process2() {
    auto bSize = r * c;
    for (int64_t bi = 0; bi < b; ++bi) {
      for (int64_t ri = 0; ri < r; ++ri) {
        for (int64_t ci = 0; ci < c; ++ci) {
          int64_t index = bi * bSize + ri * c + ci;
          if (ri == ci) {
            GTensor(Y).SetValue(index, _DT_Y(1));
          } else {
            GTensor(Y).SetValue(index, _DT_Y(0));
          }
        }
      }
    }
  }
};
} // namespace AscendC
extern "C" __global__ __aicore__ void eye(GM_ADDR y, GM_ADDR y_ref,
                                          GM_ADDR workspace, GM_ADDR tiling) {
  GET_TILING_DATA(tiling_data, tiling);
  AscendC::Kernel<DTYPE_Y> op;
  op.Init(y, y_ref, tiling_data);
  op.Process2();
}